#include <cuda_fp16.h>
#include <cuda_runtime.h>

#include <cfloat>
#include <cstdio>

#include "CudaFloats.cuh"

#define DOUBLE4(value) reinterpret_cast<double4*>(&(value))[0]

namespace cufloat
{
    //xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx atomic min max
    __device__ double atomicMin_double(double* address, double val)
    {
        // 将 double* 转换为 unsigned long long int* 以便进行位操作
        unsigned long long int* address_as_ull = (unsigned long long int*)address;
        // 读取当前存储在 address 中的值（以 unsigned long long int 形式）
        unsigned long long int old = *address_as_ull;
        unsigned long long int assumed;

        do
        {
            // 在循环开始，假设当前存储的值(old)就是我们上次比较的值
            assumed = old;
            // 计算新的候选值：当前线程的 val 和 assumed 转换回的 double 值中的较小者
            // 然后用原子操作尝试更新
            old = atomicCAS(address_as_ull, assumed, __double_as_longlong(fmin(val, __longlong_as_double(assumed))));
            // 如果 atomicCAS 时地址的值已经被其他线程改变（即 old != assumed），则循环继续尝试
        } while (assumed != old);

        // 返回原始地址中的 double 值
        return __longlong_as_double(old);
    }

    __device__ double atomicMax_double(double* address, double val)
    {
        unsigned long long int* address_as_ull = (unsigned long long int*)address;
        unsigned long long int old = *address_as_ull;
        unsigned long long int assumed;

        do
        {
            assumed = old;
            old = atomicCAS(address_as_ull, assumed, __double_as_longlong(fmax(val, __longlong_as_double(assumed))));
        } while (assumed != old);

        return __longlong_as_double(old);
    }

    //xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx

    // 自定义 double 类型的原子最小值操作
    __device__ double AtomicMinDouble(double* address, double val)
    {
        // 将 double 指针转换为 unsigned long long 指针（利用 64 位存储特性）
        unsigned long long* address_as_ull = (unsigned long long*)address;
        unsigned long long old = *address_as_ull;  // 读取当前值
        unsigned long long assumed;

        do
        {
            assumed = old;
            // 计算新值（当前最小值与输入值的较小者）
            double current_min = __longlong_as_double(assumed);
            double new_min = fmin(val, current_min);
            // 使用 atomicCAS 进行原子更新
            old = atomicCAS(address_as_ull, assumed, __double_as_longlong(new_min));
        } while (assumed != old);  // 循环直到更新成功

        return __longlong_as_double(old);
    }

    // 自定义 double 类型的原子最大值操作
    __device__ double AtomicMaxDouble(double* address, double val)
    {
        unsigned long long* address_as_ull = (unsigned long long*)address;
        unsigned long long old = *address_as_ull;
        unsigned long long assumed;

        do
        {
            assumed = old;
            double current_max = __longlong_as_double(assumed);
            double new_max = fmax(val, current_max);
            old = atomicCAS(address_as_ull, assumed, __double_as_longlong(new_max));
        } while (assumed != old);

        return __longlong_as_double(old);
    }
}  // namespace cufloat

//xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
__global__ void CuReduceMinMax(int len, double* d_x, double* d_min, double* d_max)
{
    const size_t idx = (blockDim.x * blockIdx.x + threadIdx.x) * 4;
    // for idx > N, 不做处理, 但是需要初始化 d_min, d_max;

#if 0
    // 初始化全局内存中的d_min和d_max，确保只有一个线程执行初始化
    if (threadIdx.x == 0 && blockIdx.x == 0)
    {
        *d_min = DBL_MAX;
        *d_max = -DBL_MAX;
    }
    __syncthreads();
#else
    *d_min = DBL_MAX;
    *d_max = -DBL_MAX;
#endif

    __shared__ double s_min[32];
    __shared__ double s_max[32];

#if 0
    // 初始化共享内存数组
    if (threadIdx.x < 32)
    {
        s_min[threadIdx.x] = DBL_MAX;
        s_max[threadIdx.x] = -DBL_MAX;
    }
    __syncthreads();
#endif

    const int warpId = threadIdx.x / warpSize;
    const int laneId = threadIdx.x % warpSize;

    double thread_min = DBL_MAX;
    double thread_max = -DBL_MAX;

    if (idx < len)
    {
        double4 tmp = DOUBLE4(d_x[idx]);
        //printf("[my] tmp -> %f, %f, %f, %f", tmp.x, tmp.y, tmp.z, tmp.w);
        thread_min = fmin(fmin(tmp.x, tmp.y), fmin(tmp.z, tmp.w));
        thread_max = fmax(fmax(tmp.x, tmp.y), fmax(tmp.z, tmp.w));
        //printf("[my]dbl4 idx %d, min %f, max %f", idx, thread_min, thread_max);
    }

/// Warp-level reduction for min
#pragma unroll
    for (int offset = warpSize >> 1; offset > 0; offset >>= 1)
    {
        double shuffled_min = __shfl_down_sync(0xFFFFFFFF, thread_min, offset);
#if 0
        if (shuffled_min == 0.0)
        {
            printf("[my] warp %d, val->%f", warpId, shuffled_min);
        }
#endif
        thread_min = fmin(thread_min, shuffled_min);
    }

/// Warp-level reduction for max
#pragma unroll
    for (int offset = warpSize >> 1; offset > 0; offset >>= 1)
    {
        double shuffled_max = __shfl_down_sync(0xFFFFFFFF, thread_max, offset);
        thread_max = fmax(thread_max, shuffled_max);
    }

    if (laneId == 0)
    {
        s_min[warpId] = thread_min;
        s_max[warpId] = thread_max;
        // printf("[my]warpId %d; red, min %f, max %f", warpId, s_min[warpId], s_max[warpId]);
    }

    __syncthreads();

    // First warp reduces the results from all warps
    if (warpId == 0)
    {
        double block_min = (laneId < blockDim.x / warpSize) ? s_min[laneId] : DBL_MAX;
        double block_max = (laneId < blockDim.x / warpSize) ? s_max[laneId] : -DBL_MAX;

#pragma unroll
        for (int offset = warpSize / 2; offset > 0; offset /= 2)
        {
            double shuffled_min = __shfl_down_sync(0xFFFFFFFF, block_min, offset);
            double shuffled_max = __shfl_down_sync(0xFFFFFFFF, block_max, offset);
            block_min = fmin(block_min, shuffled_min);
            block_max = fmax(block_max, shuffled_max);
        }

        if (laneId == 0)
        {
            cufloat::atomicMin_double(d_min, block_min);
            cufloat::atomicMax_double(d_max, block_max);
#if 0
            printf("[my lane 0] min %f max %f\n", d_min, d_max);
#endif
        }
    }
}
